// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package org.jetbrains.kotlin.idea.k2.codeinsight.hierarchy.calls

import com.intellij.ide.hierarchy.HierarchyNodeDescriptor
import com.intellij.ide.hierarchy.call.CallHierarchyNodeDescriptor
import com.intellij.ide.util.treeView.NodeDescriptor
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiMember
import com.intellij.psi.PsiReference
import org.jetbrains.kotlin.asJava.toLightElements
import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf

fun isCallHierarchyElement(e: PsiElement): Boolean {
    return (e is KtNamedFunction && e.name != null) ||
            e is KtSecondaryConstructor ||
            (e is KtProperty && !e.isLocal) ||
            e is KtObjectDeclaration ||
            (e is KtClass && !e.isInterface()) ||
            e is KtFile
}

fun getCallHierarchyElement(element: PsiElement) = element.parentsWithSelf.firstOrNull(::isCallHierarchyElement) as? KtElement

private fun NodeDescriptor<*>.incrementUsageCount() {
    when (this) {
        is KotlinCallHierarchyNodeDescriptor -> incrementUsageCount()
        is CallHierarchyNodeDescriptor -> incrementUsageCount()
    }
}

private fun NodeDescriptor<*>.addReference(reference: PsiReference) {
    when (this) {
        is KotlinCallHierarchyNodeDescriptor -> addReference(reference)
        is CallHierarchyNodeDescriptor -> addReference(reference)
    }
}

internal fun getOrCreateNodeDescriptor(
    parent: HierarchyNodeDescriptor,
    originalElement: PsiElement,
    reference: PsiReference?,
    navigateToReference: Boolean,
    elementToDescriptorMap: MutableMap<PsiElement, NodeDescriptor<*>>,
    isJavaMap: Boolean
): HierarchyNodeDescriptor? {
    val element = (if (isJavaMap && originalElement is KtElement) originalElement.toLightElements().firstOrNull() else originalElement)
        ?: return null

    val existingDescriptor = elementToDescriptorMap[element] as? HierarchyNodeDescriptor
    val result = if (existingDescriptor != null) {
        existingDescriptor.incrementUsageCount()
        existingDescriptor
    } else {
        val newDescriptor: HierarchyNodeDescriptor = when (element) {
            is KtElement -> KotlinCallHierarchyNodeDescriptor(parent, element, false, navigateToReference)
            is PsiMember -> CallHierarchyNodeDescriptor(element.project, parent, element, false, navigateToReference)
            else -> return null
        }
        elementToDescriptorMap[element] = newDescriptor
        newDescriptor
    }

    if (reference != null) {
        result.addReference(reference)
    }

    return result
}